# -*- coding: UTF-8 -*-
"""
@项目名称：pdb_add_h.py
@作   者：陆地起飞全靠浪
@创建日期：2025-09-19-10:43
"""
from openbabel import pybel
from rdkit import Chem


def ligand_add_hydrogens(input_ligand_pdb, output_ligand_addh):
    mol = Chem.MolFromPDBFile(input_ligand_pdb)
    # mol = Chem.RemoveHs(mol)
    mol = Chem.AddHs(mol)  # 添加氢原子
    Chem.MolToPDBFile(mol, output_ligand_addh)
    print(f'删除旧氢，添加新氢完成，结果保存位置为：{output_ligand_addh}')


def add_hydrogens(input_ligand_pdb, output_protein_addh):
    """
    处理PDB文件中的氢原子：先移除所有现有氢原子，然后重新添加氢原子
    参数:
        input_ligand_pdb (str): 输入的PDB文件路径
        output_protein_addh (str): 输出的PDB文件路径
    """
    try:
        # 读取PDB文件
        molecules = list(pybel.readfile("pdb", input_ligand_pdb))
        if not molecules:
            print(f"错误: 文件 '{input_ligand_pdb}' 中没有读取到分子")
            return False
        # 处理每个分子（PDB文件通常包含一个分子，但可能有多个模型）
        processed_molecules = []
        for mol in molecules:
            # 1. 移除所有现有的氢原子
            mol.removeh()
            print("已移除所有现有氢原子")
            # 2. 添加新的氢原子
            mol.addh()
            print("已添加新的氢原子")
            processed_molecules.append(mol)
        # 写入输出文件
        with open(output_protein_addh, 'w+') as out_file:
            for mol in processed_molecules:
                out_file.write(mol.write("pdb"))
        print(f"处理完成！输出文件已保存至: {output_protein_addh}")
        return True
    except Exception as e:
        print(f"处理过程中发生错误: {str(e)}")
        return False


# 示例使用
if __name__ == "__main__":
    input_file = "/4T/AIDDserver/MolecularDocking/DockResult/docking_2025_09_22-08_32_42/4bvw/chains/chain_配体_C.pdb"  # 请替换为你的输入PDB文件路径
    output_file = "/4T/AIDDserver/MolecularDocking/DockResult/docking_2025_09_22-08_32_42/4bvw/4bvw_ligand_addh1.pdb"  # 输出PDB文件路径
    # 处理PDB文件
    ligand_add_hydrogens(input_file, output_file)
    input_file = "/4T/AIDDserver/MolecularDocking/DockResult/docking_2025_09_22-08_32_42/4bvw/chains/chain_配体_C.pdb"  # 请替换为你的输入PDB文件路径
    output_file = "/4T/AIDDserver/MolecularDocking/DockResult/docking_2025_09_22-08_32_42/4bvw/4bvw_ligand_addh2.pdb"  # 输出PDB文件路径
    add_hydrogens(input_file, output_file)
